package com.acc.accrpc.loadbalancer;

import com.acc.accrpc.model.ServiceMetaInfo;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

/**
 * 一致性哈希负载均衡器
 */
public class ConsistentHashLoadBalancer implements LoadBalancer {

    /**
     * 一致性 Hash 环，存放虚拟节点
     */
    private final TreeMap<Long, ServiceMetaInfo> virtualNodes = new TreeMap<>();

    /**
     * 虚拟节点数
     */
    private static final int VIRTUAL_NODE_NUM = 100;

    @Override
    public ServiceMetaInfo select(Map<String, Object> requestParams, List<ServiceMetaInfo> serviceMetaInfoList) {
        if (serviceMetaInfoList.isEmpty()) {
            return null;
        }

        // 清空旧节点，放置节点堆积
        virtualNodes.clear();

        // 构建虚拟节点环
        for (ServiceMetaInfo serviceMetaInfo : serviceMetaInfoList) {
            for (int i = 0; i < VIRTUAL_NODE_NUM; i++) {
                // 每个真实地址生成多个虚拟地址
                // 127.0.0.1:80 -> 127.0.0.1:800
                // 127.0.0.1:801、127.0.0.1:802...
                String virtualAddress = serviceMetaInfo.getServiceAddress() + i;
                // 每个虚拟地址获取4个虚拟节点
                for (int j = 0; j < 4; j++) {
                    long hash = getHash(virtualAddress, j);
                    virtualNodes.put(hash, serviceMetaInfo);
                }
            }
        }

        // 获取调用请求的 hash 值
        // 取 md5 摘要最后一部分作为 hash
        long hash = getHash(requestParams.get("client_address"), 3);
        // 选择最接近且大于等于调用请求 hash 值的虚拟节点
        Map.Entry<Long, ServiceMetaInfo> entry = virtualNodes.ceilingEntry(hash);
        if (entry == null) {
            // 如果没有大于等于调用请求 hash 值的虚拟节点，则返回环首部的节点
            entry = virtualNodes.firstEntry();
        }
        return entry.getValue();
    }


    /**
     * Hash 算法
     * md5哈希算法实现
     *
     * @param key
     * @return
     */
    private long getHash(Object key, int number) {
        try {
            /*
                利用 md5 生成的 16 字节摘要，生成四个 32 位哈希值,
                每个哈希值相当于一个虚拟节点
             */
            MessageDigest md = MessageDigest.getInstance("MD5");
            byte[] digest = md.digest(key.toString().getBytes(StandardCharsets.UTF_8));

            /*
                将 digest 字节数组划分为 4 个部分，每个部分 4字节
                每次调用返回 number 指定部分的 32 位整数作为哈希值
             */
            return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                    | (digest[number * 4] & 0xFF))
                    & 0xFFFFFFFFL;
        } catch (NoSuchAlgorithmException e) {
            // 理论上不会发生，MD5算法在所有Java实现中都应该存在
            throw new RuntimeException("MD5 algorithm not available", e);
        }
    }
}
